from transformers import AutoTokenizer, AutoModelForSeq2SeqLM

import time

device = 'cpu' #or 'cpu' for translate on cpu

model = AutoModelForSeq2SeqLM.from_pretrained("./models/opus-mt-zh-en")
model.to(device)

tokenizer = AutoTokenizer.from_pretrained("./models/opus-mt-zh-en")

def getModelAndTokenizer():
    return model, tokenizer

if __name__ == "__main__":
    src_text = "基于知行合一理念，高校思政教育教学工作者应恪守示范性原则，切实做到为人师表、以身作则，真正做到理论和实践的统一，以榜样的力量感染和带动学生，使学生在教育过程中受到感染和启发，并以教师为标榜，严格要求、约束和激励自我，最终内化为良好的知行合一品质。此外，教师应发挥教学与实践的主导作用，切实履行好自身的职责与使命，善于因势利导、因时而新，以深厚的理论功底赢得学生，提升育人的亲和度。"

    start = time.time()
    input_ids = tokenizer(src_text, return_tensors="pt")

    generated_tokens = model.generate(**input_ids.to(device))

    result = tokenizer.batch_decode(generated_tokens, skip_special_tokens=True)
    end = time.time()
    print(end-start,' s')
    print(result)